# Copyright (C) 2015-2024, Wazuh Inc.
# Created by Wazuh, Inc. <info@wazuh.com>.
# This program is free software; you can redistribute it and/or modify it under the terms of GPLv2
import pytest
import os
from pathlib import Path

from wazuh_testing.utils.db_queries import cve_db
from wazuh_testing.utils.services import control_service
from wazuh_testing import FEEDS_PATH
from wazuh_testing.constants.paths.databases import CPE_HELPER_PATH
from wazuh_testing.modules.modulesd.vulnerability_detector import CUSTOM_CPE_HELPER
from wazuh_testing.utils.file import read_json_file, copy, write_json_file
from wazuh_testing.utils.mocking import (insert_vulnerable_packages,
                                         insert_vulnerabilities_agent_inventory,
                                         insert_suse_system_package, VULNERABLE_PACKAGES)
from wazuh_testing.utils.db_queries import agent_db
from wazuh_testing.utils.time import get_current_timestamp


CUSTOM_CPE_HELPER_PATH = str(Path(FEEDS_PATH, 'cpe_helper', CUSTOM_CPE_HELPER))


@pytest.fixture()
def clean_cve_tables() -> None:
    """Clean all tables of the CVE database before and after finishing the test"""
    cve_db.clean_all_cve_tables()
    yield
    cve_db.clean_all_cve_tables()


@pytest.fixture(scope='module')
def clean_cve_tables_module() -> None:
    """Clean all tables of the CVE database before and after finishing the test"""
    cve_db.clean_all_cve_tables()
    yield
    cve_db.clean_all_cve_tables()


@pytest.fixture()
def prepare_full_scan(agent_system, mock_agent_with_custom_system):
    """Prepare the environment to launch the vulnerability scan.

    - Mock an agent with a specified system.
    - Force full scan.

    Args:
        agent_system (str): System to set to the mocked agent.
        mock_agent_with_custom_system (fixture): Mock an agent with a custom system.
    """
    agent_db.update_last_full_scan(1, agent_id=mock_agent_with_custom_system)

    yield mock_agent_with_custom_system


def prepare_scan_with_vuln_packages_implementation(agent_id: int, last_scan: int, agent_system: str = 'RHEL8',
                                                   count: int = 5) -> None:
    """Prepare the environment to launch the vulnerability scan.

    - Mock an agent with a specified system.
    - Insert mocked vulnerables packages.
    - Update packages sync status.
    - Force full scan.

    Args:
        agent_id (int): Agent id to which packages and scan will be configured.
        last_scan (int): Timestamp to set as last scan. used to force different types of scans
        agent_system (str): System to set to the mocked agent.
        count (int): Number of packages to insert.
    """
    # Insert vulnerable packages
    package_vendor = 'Red Hat, Inc.' if 'RHEL' in agent_system else 'wazuh-mocking'

    if 'SLES' in agent_system:
        package_vendor = 'SUSE LLC <https://www.suse.com/>'
        insert_suse_system_package(agent_id=agent_id, version=agent_system)

    insert_vulnerable_packages(agent_id=agent_id, vendor=package_vendor, count=count)

    # Update sync info for packages
    agent_db.update_sync_info(agent_id=agent_id, component="syscollector-packages")

    # Forcing a scan
    agent_db.update_last_full_scan(last_scan, agent_id=agent_id)


@pytest.fixture()
def prepare_baseline_scan_with_vuln_packages(mock_agent_with_custom_system: int, agent_system: str) -> int:
    """Add a mocked agent with mocked packages and force the baseline scan for that agent.

    Args:
        mock_agent_with_custom_system (Fixture): Fixture for mocking an agent.
        agent_system (str): System to set to the mocked agent.
    """
    prepare_scan_with_vuln_packages_implementation(agent_id=mock_agent_with_custom_system, last_scan=0,
                                                   agent_system=agent_system)

    yield mock_agent_with_custom_system


@pytest.fixture()
def prepare_full_scan_with_vuln_packages(mock_agent_with_custom_system: int, agent_system: str) -> int:
    """Prepare the environment to launch the vulnerability full scan.

    Args:
        mock_agent_with_custom_system (fixture): Mock an agent with a custom system.
        agent_system (str): System to set to the mocked agent.
    """
    prepare_scan_with_vuln_packages_implementation(agent_id=mock_agent_with_custom_system, last_scan=1,
                                                   agent_system=agent_system)

    yield mock_agent_with_custom_system


@pytest.fixture()
def prepare_full_scan_with_vuln_package(mock_agent_with_custom_system: int, agent_system: str) -> int:
    """Add a mocked agent with only one mocked vulnerable package and force the full scan for that agent.

    Args:
        mock_agent_with_custom_system (fixture): Mock an agent with a custom system.
        agent_system (str): System to set to the mocked agent.
    """
    prepare_scan_with_vuln_packages_implementation(agent_id=mock_agent_with_custom_system, last_scan=1,
                                                   agent_system=agent_system, count=1)

    yield mock_agent_with_custom_system


@pytest.fixture()
def prepare_partial_scan_with_vuln_packages(mock_agent_with_custom_system, agent_system: str) -> int:
    """Add a mocked agent with mocked packages and force the partial scan for that agent.

    Args:
        mock_agent_with_custom_system (Fixture): Fixture for mocking an agent.
        agent_system (str): System to set to the mocked agent.
    """
    prepare_scan_with_vuln_packages_implementation(agent_id=mock_agent_with_custom_system,
                                                   last_scan=int(get_current_timestamp()), agent_system=agent_system)

    yield mock_agent_with_custom_system


def prepare_scan_with_obsolete_vulnerabilities_implementation(agent_id: int, last_scan: int,
                                                              agent_system: str = 'RHEL8') -> None:
    """Prepare the environment to launch the vulnerability scan.

    - Insert mocked vulnerables packages.
    - Update packages sync status.
    - Force scan based on the last_scan time passed.

    Args:
        agent_id (int): Agent id to insert the packages.
        last_scan (int): timestamp of last previous scan.
        agent_system (str): System to set to the mocked agent. Default: 'RHEL8'
    """
    # Insert vulnerability data in vulnerabilities inventory
    insert_vulnerabilities_agent_inventory(agent_id=agent_id, status='OBSOLETE')

    # Force sync status for packages
    agent_db.update_sync_info(agent_id=agent_id, component="syscollector-packages")

    # Force the partial scan
    agent_db.update_last_full_scan(last_scan, agent_id=agent_id)


@pytest.fixture()
def prepare_partial_scan_with_obsolete_vulnerabilities(mock_agent_with_custom_system: int,
                                                       agent_system: str) -> int:
    """Add a mocked agent with mocked packages and force the partial scan for that agent.

    Args:
        mock_agent_with_custom_system (Fixture): Fixture for mocking an agent.
        agent_system (str): System to set to the mocked agent.
    """
    prepare_scan_with_obsolete_vulnerabilities_implementation(agent_id=mock_agent_with_custom_system,
                                                              last_scan=int(get_current_timestamp()),
                                                              agent_system=agent_system)

    yield mock_agent_with_custom_system


@pytest.fixture()
def prepare_full_scan_with_obsolete_vulnerabilities(mock_agent_with_custom_system: int,
                                                    agent_system: str) -> int:
    """Add a mocked agent with mocked packages and force the full scan for that agent.

    Args:
        mock_agent_with_custom_system (Fixture): Fixture for mocking an agent.
        agent_system (str): System to set to the mocked agent.
    """
    prepare_scan_with_obsolete_vulnerabilities_implementation(agent_id=mock_agent_with_custom_system,
                                                              last_scan=1, agent_system=agent_system)

    yield mock_agent_with_custom_system


@pytest.fixture()
def prepare_scan_with_cpe_helper(agent_system, mock_agent_with_custom_system):
    """Prepare the environment to launch the vulnerability scan.

    - Mock an agent with a specified system.
    - Insert mocked vulnerables packages.
    - Update packages sync status.
    - Copy the custom CPE helper to the dictionaries folder.

    Args:
        agent_system (str): System to set to the mocked agent.
        mock_agent_with_custom_system (fixture): Mock an agent with a custom system.
    """
    for package in VULNERABLE_PACKAGES:
        agent_db.insert_package(name=package['name'], version=package['version'], source=package['name'],
                                vendor=package['vendor'], agent_id=mock_agent_with_custom_system)

    # Sync packages info
    agent_db.update_sync_info(agent_id=mock_agent_with_custom_system, component="syscollector-packages")
    agent_db.update_sync_info(agent_id=mock_agent_with_custom_system, component="syscollector-hotfixes")

    # Make a backup data from inital CPE helper
    cpe_helper_backup_data = read_json_file(CPE_HELPER_PATH)

    # Set the custom CPE helper
    copy(CUSTOM_CPE_HELPER_PATH, CPE_HELPER_PATH)

    yield mock_agent_with_custom_system

    # Restore the CPE helper backup data
    write_json_file(CPE_HELPER_PATH, cpe_helper_backup_data)


@pytest.fixture()
def prepare_full_scan_environment(mock_agent_with_custom_system, mock_agent_packages):
    """Add a mocked agent with mocked packages and force the full scan for that agent.

    Args:
        mock_agent_with_custom_system (Fixture): Fixture for mocking an agent.
        mock_agent_packages (Fixture): Add mocked packages for mocked agent.
    """
    # Force sync status for packages
    print("Updating sync info for agent: {}".format(mock_agent_with_custom_system))
    print("Updating sync info for agent: {}".format(mock_agent_packages))
    agent_db.update_sync_info(agent_id=mock_agent_with_custom_system, component="syscollector-packages")

    # Force the full scan
    agent_db.update_last_full_scan(1, agent_id=mock_agent_with_custom_system)

    yield mock_agent_with_custom_system


@pytest.fixture()
def prepare_scan(metadata, agent_system, prepare_full_scan_with_vuln_packages):
    """Prepare the environment to launch the vulnerability scan.

    - Mock an agent with a specified system.
    - Insert mocked vulnerables packages.
    - Update packages sync status.
    - If Windows, then copy the custom CPE helper to the dictionaries folder.
    - Force full scan.

    Args:
        metadata (dict): Test case metadata.
        agent_system (str): System to set to the mocked agent.
        prepare_full_scan_with_vuln_packages (fixture): Mock an agent with a custom system and vulnerable packages.
    """

    if 'WINDOWS' in agent_system:
        agent_db.update_sync_info(agent_id=prepare_full_scan_with_vuln_packages, component="syscollector-hotfixes")
        # Make a backup data from inital CPE helper
        cpe_helper_backup_data = read_json_file(CPE_HELPER_PATH)
        # Set the custom CPE helper
        copy(CUSTOM_CPE_HELPER_PATH, CPE_HELPER_PATH)

    # Forcing a full-scan
    agent_db.update_last_full_scan(1, agent_id=prepare_full_scan_with_vuln_packages)

    yield prepare_full_scan_with_vuln_packages

    if 'WINDOWS' in agent_system:
        # Restore the CPE helper backup data
        write_json_file(CPE_HELPER_PATH, cpe_helper_backup_data)
